#include"fbs.hpp"

#if __cplusplus <= 199711L
    #include <map>
    typedef std::map<long long /* hash */, int /* vert id */>  mapId;
#else
    #include <unordered_map>
    typedef std::unordered_map<long long /* hash */, int /* vert id */>  mapId;
#endif


void r3d_depth_image::FastBilateralSolverFilter_::init(
        cv::Mat& reference,
		cv::Mat& depth,
        double sigma_spatial,
        double sigma_luma,
        double sigma_chroma,
		double sigma_depth,
        double lambda,
        int max_iter,
        double max_tol,
        const cv::Mat & process_mask,
        int stride,
        const cv::Mat & edge_mask
) {
	IMAGE_DEPTH_LOGI("FastBilateralSolverFilter_::init() begin");
    double start, stop, freq = cv::getTickFrequency();

    bs_param.lam = lambda;
    bs_param.cg_maxiter = max_iter;
    bs_param.cg_tol = max_tol;

    if(reference.channels()==1) {
		IMAGE_DEPTH_LOGI("[FastBilateralSolverFilter_]: guide image has to be of 3 channel!");
        exit(-1);
    } else {
        const int dim = 6;
        cv::Mat reference_yuv;
        cv::cvtColor(reference, reference_yuv, COLOR_BGR2YCrCb);

        cols = reference_yuv.cols;
        rows = reference_yuv.rows;
        npixels = rows * cols;

        long long hash_vec[dim];
        for (int i = 0; i < dim; ++i) hash_vec[i] = static_cast<long long>(std::pow(255, i));

        mapId hashed_coords;
#if __cplusplus <= 199711L
#else
        hashed_coords.reserve(cols*rows);
#endif

        const unsigned char* pRef = reference_yuv.data;
        const unsigned char* pMask = process_mask.data;
        const unsigned char* pEdge = edge_mask.data;
		const unsigned char* pDepth = depth.data;
        int vert_idx = 0;

        // construct Splat(Slice) matrices
        splat_idx.resize(npixels);
        full_splat_idx.resize(npixels);

        start = cv::getTickCount();
        // initialize the splat index map
        int * pSplatIdx = splat_idx.data();
        int * pFullSplatIdx = full_splat_idx.data();
        size_t sizeSplatIdx = splat_idx.size();
        for(int i=0; i<sizeSplatIdx; ++i) {
            pSplatIdx[i] = -1;
            pFullSplatIdx[i] = -1;
        }
        stop = cv::getTickCount();
		IMAGE_DEPTH_LOGI("init splat index: %.3f ms", 1000*(stop-start)/freq);

        start = cv::getTickCount();
        for (int y = 0; y < rows; ++y) {
            for (int x = 0; x < cols; ++x) {
                int pix_idx = y * cols + x;
                if ( pMask[pix_idx] ) {
                    int ref_idx = pix_idx * 3;
                    long long coord[dim];
                    coord[0] = int(x / sigma_spatial); // X COORD
                    coord[1] = int(y / sigma_spatial); // Y COORD
                    coord[2] = int(pRef[ref_idx] / sigma_luma); // L COORD
                    coord[3] = int(pRef[ref_idx + 1] / sigma_chroma); // U COORD
                    coord[4] = int(pRef[ref_idx + 2] / sigma_chroma); // V COORD
					coord[5] = int(pDepth[pix_idx] / sigma_depth); // D coord

                    // convert the coordinate to a hash value
                    long long hash_coord = 0;
                    for (int i = 0; i < dim; ++i) hash_coord += coord[i] * hash_vec[i];

                    // pixels whom are alike will have the same hash value.
                    // We only want to keep a unique list of hash values,
                    // therefore make sure we only insert unique hash values.
                    //if ( pEdge[pix_idx] || (x % stride == 0 && y % stride == 0) ) {
                    if ( x % stride == 0 && y % stride == 0 ) {
                        mapId::iterator it = hashed_coords.find(hash_coord);
                        if (it == hashed_coords.end()) {
                            hashed_coords.insert(
                                    std::pair<long long, int>(hash_coord, vert_idx));
                            splat_idx[pix_idx] = vert_idx;
                            full_splat_idx[pix_idx] = vert_idx;
                            ++vert_idx;
                        } else {
                            splat_idx[pix_idx] = it->second;
                            full_splat_idx[pix_idx] = it->second;
                        }
                    } else {
                        mapId::iterator it = hashed_coords.find(hash_coord);
                        if (it == hashed_coords.end()) {
                            full_splat_idx[pix_idx] = -2;
                        } else {
                            full_splat_idx[pix_idx] = it->second;
                        }
                    }
                } else {
                    int ref_idx = pix_idx * 3;
                    long long coord[dim];
                    coord[0] = int(x / sigma_spatial); // X COORD
                    coord[1] = int(y / sigma_spatial); // Y COORD
                    coord[2] = int(pRef[ref_idx] / sigma_luma); // L COORD
                    coord[3] = int(pRef[ref_idx + 1] / sigma_chroma); // U COORD
                    coord[4] = int(pRef[ref_idx + 2] / sigma_chroma); // V COORD
					coord[5] = int(pDepth[pix_idx] / sigma_depth); // D coord

                    // convert the coordinate to a hash value
                    long long hash_coord = 0;
                    for (int i = 0; i < dim; ++i) hash_coord += coord[i] * hash_vec[i];

                    // pixels whom are alike will have the same hash value.
                    // We only want to keep a unique list of hash values, therefore make sure we only insert
                    // unique hash values.
                    if (x % (stride<<1) == 0 && y % (stride<<1) == 0) {
                        mapId::iterator it = hashed_coords.find(hash_coord);
                        if (it == hashed_coords.end()) {
                            hashed_coords.insert(
                                    std::pair<long long, int>(hash_coord, vert_idx));
                            splat_idx[pix_idx] = vert_idx;
                            full_splat_idx[pix_idx] = vert_idx;
                            ++vert_idx;
                        } else {
                            splat_idx[pix_idx] = it->second;
                            full_splat_idx[pix_idx] = it->second;
                        }
                    } else {
                        mapId::iterator it = hashed_coords.find(hash_coord);
                        if (it == hashed_coords.end()) {
                            full_splat_idx[pix_idx] = -2;
                        } else {
                            full_splat_idx[pix_idx] = it->second;
                        }
                    }
                }
            }
        }
        nvertices = static_cast<int>(hashed_coords.size());
		IMAGE_DEPTH_LOGI("vertices: %d", nvertices);
        stop = cv::getTickCount();
		IMAGE_DEPTH_LOGI("generate hash table: %.3f ms", 1000*(stop-start)/freq);

        // construct Blur matrices
        start = cv::getTickCount();
        Eigen::VectorXf ones_nvertices = Eigen::VectorXf::Ones(nvertices);
        diagonal(ones_nvertices, blurs);
        blurs *= 10;
        for(int offset = -1; offset <= 1; ++offset) {
            if(offset == 0) continue;
            for (int i = 0; i < dim; ++i) {
                Eigen::SparseMatrix<float, Eigen::ColMajor> \
blur_temp(hashed_coords.size(), hashed_coords.size());
                blur_temp.reserve(Eigen::VectorXi::Constant(nvertices, 6));
                long long offset_hash_coord = offset * hash_vec[i];
                for (mapId::iterator it = hashed_coords.begin(); \
                it != hashed_coords.end(); ++it) {
                    long long neighb_coord = it->first + offset_hash_coord;
                    mapId::iterator it_neighb = hashed_coords.find(neighb_coord);
                    if (it_neighb != hashed_coords.end()) {
                        blur_temp.insert(it->second,it_neighb->second) = 1.0f;
                        blur_idx.push_back(std::pair<int,int>(it->second, it_neighb->second));
                    }
                }
                blurs += blur_temp;
            }
        }
        blurs.finalize();
        stop = cv::getTickCount();
		IMAGE_DEPTH_LOGI("generate blur matrices: %.3f ms", 1000*(stop-start)/freq);

        //bistochastize
        start = cv::getTickCount();
        int maxiter = 10;
        n = ones_nvertices;
        m = Eigen::VectorXf::Zero(nvertices);
        for (int i = 0; i < int(splat_idx.size()); i++) {
            if (splat_idx[i] >= 0) m(splat_idx[i]) += 1.0f;
        }

        Eigen::VectorXf bluredn(nvertices);

        for (int i = 0; i < maxiter; i++) {
            Blur(n, bluredn);
            n = ((n.array()*m.array()).array()/bluredn.array()).array().sqrt();
        }
        Blur(n, bluredn);

        m = n.array() * (bluredn).array();
        diagonal(m, Dm);
        diagonal(n, Dn);
        stop = cv::getTickCount();
		IMAGE_DEPTH_LOGI("bistochastize: %.3f ms", 1000*(stop-start)/freq);
    }
	IMAGE_DEPTH_LOGI("FastBilateralSolverFilter_::init() end");
}

void r3d_depth_image::FastBilateralSolverFilter_::Splat(Eigen::VectorXf& input, Eigen::VectorXf& output) {
    output.setZero();
    for (int i = 0, t_=0; i < int(splat_idx.size()); i++) {
        if ( splat_idx[i] >= 0 ) output(splat_idx[i]) += input(t_++);
    }
}

void r3d_depth_image::FastBilateralSolverFilter_::Blur(Eigen::VectorXf& input, Eigen::VectorXf& output) {
    output.setZero();
    output = input * 10;
    for (int i = 0; i < int(blur_idx.size()); i++) {
        output(blur_idx[i].first) += input(blur_idx[i].second);
    }
}

void r3d_depth_image::FastBilateralSolverFilter_::solve(
        cv::Mat& target,
        cv::Mat& confidence,
        cv::Mat& output) {
	IMAGE_DEPTH_LOGI("FastBilateralSolverFilter_::solve() begin");
    double start, stop, freq = cv::getTickFrequency();

    Eigen::SparseMatrix<float, Eigen::ColMajor> M(nvertices, nvertices);
    Eigen::SparseMatrix<float, Eigen::ColMajor> A_data(nvertices, nvertices);
    Eigen::SparseMatrix<float, Eigen::ColMajor> A(nvertices, nvertices);
    Eigen::VectorXf b(nvertices);
    Eigen::VectorXf y(nvertices);
    Eigen::VectorXf y0(nvertices);
    Eigen::VectorXf y1(nvertices);
    Eigen::VectorXf w_splat(nvertices);

    // construct input and weight matrix
    start = cv::getTickCount();
    num_eq = 0; // number of equations to solve
    for (int i=0; i<npixels; ++i) if (splat_idx[i]>=0) num_eq++;
    Eigen::VectorXf x(num_eq);
    Eigen::VectorXf w(num_eq);

    if(target.depth() == CV_16S) {
        const int16_t *pft = reinterpret_cast<const int16_t*>(target.data);
        for (int i = 0, t_=0; i < npixels; i++) {
            if ( splat_idx[i] >= 0)
                x(t_++) = (cv::saturate_cast<float>(pft[i]) + 32768.0f) / 65535.0f;
        }
    } else if(target.depth() == CV_16U) {
        const uint16_t *pft = reinterpret_cast<const uint16_t*>(target.data);
        for (int i = 0, t_=0; i < npixels; i++)
            if ( splat_idx[i] >= 0) x(t_++) = cv::saturate_cast<float>(pft[i])/65535.0f;
    } else if(target.depth() == CV_8U) {
        const uchar *pft = reinterpret_cast<const uchar*>(target.data);
        for (int i = 0, t_=0; i < npixels; i++)
            if ( splat_idx[i] >= 0) x(t_++) = cv::saturate_cast<float>(pft[i])/255.0f;
    } else if(target.depth() == CV_32F) {
        const float *pft = reinterpret_cast<const float*>(target.data);
        for (int i = 0, t_=0; i < npixels; i++)
            if ( splat_idx[i] >= 0) x(t_++) = pft[i];
    }

    if(confidence.depth() == CV_8U) {
        const uchar *pfc = reinterpret_cast<const uchar*>(confidence.data);
        for (int i = 0, t_=0; i < npixels; i++)
            if ( splat_idx[i] >= 0) w(t_++) = cv::saturate_cast<float>(pfc[i])/255.0f;
    } else if(confidence.depth() == CV_32F) {
        const float *pfc = reinterpret_cast<const float*>(confidence.data);
        for (int i = 0, t_=0; i < npixels; i++)
            if ( splat_idx[i] >= 0) w(t_++) = pfc[i];
    }
    stop = cv::getTickCount();
	IMAGE_DEPTH_LOGI("process target and confidence map: %.3f ms", 1000*(stop - start)/freq);

    //construct A
    start = cv::getTickCount();
    Splat(w, w_splat);
    stop = cv::getTickCount();
	IMAGE_DEPTH_LOGI("Splat: %.3f ms", 1000*(stop - start)/freq);
    start = cv::getTickCount();
    diagonal(w_splat, A_data);
    stop = cv::getTickCount();
	IMAGE_DEPTH_LOGI("diagonal: %.3f ms", 1000*(stop - start)/freq);
    start = cv::getTickCount();
    A = bs_param.lam * (Dm - Dn * (blurs*Dn)) + A_data;
    stop = cv::getTickCount();
	IMAGE_DEPTH_LOGI("compute matrix A: %.3f ms", 1000*(stop - start)/freq);

    //construct b
    start = cv::getTickCount();
    b.setZero();
    for (int i = 0, t_=0; i < int(splat_idx.size()); i++) {
        if (splat_idx[i] >= 0) {
            b(splat_idx[i]) += x(t_) * w(t_);
            t_++;
        }
    }
    stop = cv::getTickCount();
	IMAGE_DEPTH_LOGI("compute vector b: %.3f ms", 1000*(stop - start)/freq);

    //construct guess for y
    start = cv::getTickCount();
    y0.setZero();
    for (int i = 0, t_=0; i < int(splat_idx.size()); i++) {
        if (splat_idx[i] >= 0) y0(splat_idx[i]) += x(t_++);
    }
    y1.setZero();
    for (int i = 0; i < int(splat_idx.size()); i++) {
        if (splat_idx[i] >= 0) y1(splat_idx[i]) += 1.0f;
    }
    for (int i = 0; i < nvertices; i++) {
        y0(i) = y0(i)/y1(i);
    }
    stop = cv::getTickCount();
	IMAGE_DEPTH_LOGI("compute guess for y: %.3f ms", 1000*(stop - start)/freq);

    // solve Ay = b
    start = cv::getTickCount();
    Eigen::ConjugateGradient<Eigen::SparseMatrix<float>, Eigen::Lower|Eigen::Upper> cg;
    cg.compute(A);
    cg.setMaxIterations(bs_param.cg_maxiter);
    cg.setTolerance(bs_param.cg_tol);
    // y = cg.solve(b);
    y = cg.solveWithGuess(b, y0);
	IMAGE_DEPTH_LOGI("Iterations:\t%d", (int)cg.iterations());
	IMAGE_DEPTH_LOGI("Loss:\t%f", (double)(cg.error()));
    stop = cv::getTickCount();
	IMAGE_DEPTH_LOGI("solve Ay=b: %.3f ms", 1000*(stop - start)/freq);

    //slice
    start = cv::getTickCount();
    if(target.depth() == CV_16S) {
        int16_t *pftar = (int16_t*) output.data;
        for (int i = 0; i < int(full_splat_idx.size()); i++) {
            if (full_splat_idx[i] >= 0)
                pftar[i] = cv::saturate_cast<short>(y(full_splat_idx[i]) * 65535.0f - 32768.0f);
        }
    } else if(target.depth() == CV_16U) {
        uint16_t *pftar = (uint16_t*) output.data;
        for (int i = 0; i < int(full_splat_idx.size()); i++) {
            if (full_splat_idx[i] >= 0)
                pftar[i] = cv::saturate_cast<ushort>(y(full_splat_idx[i]) * 65535.0f);
        }
    } else if (target.depth() == CV_8U) {
        uchar *pftar = output.data;
        for (int i = 0; i < int(full_splat_idx.size()); i++) {
            if (full_splat_idx[i] >= 0)
                pftar[i] = cv::saturate_cast<uchar>(y(full_splat_idx[i]) * 255.0f);
            else if (full_splat_idx[i] == -2)
                pftar[i] = cv::saturate_cast<uchar>(0);
        }
    } else {
        float *pftar = (float*)(output.data);
        for (int i = 0; i < int(full_splat_idx.size()); i++) {
            if (full_splat_idx[i] >= 0) pftar[i] = y(full_splat_idx[i]);
        }
    }
    stop = cv::getTickCount();
	IMAGE_DEPTH_LOGI("Slice: %.3f ms", 1000*(stop - start)/freq);

	IMAGE_DEPTH_LOGI("FastBilateralSolverFilter_::solve() end");
}

